"""
YOLOv8 数据预处理脚本
功能：
1. 检查原始图片和标注文件的匹配性
2. 自动划分训练集/验证集（默认比例 80%训练，20%验证）
3. 生成符合YOLOv8要求的数据集结构
4. 创建dataset.yaml配置文件
"""
import logging
import os
import shutil
import sys
from pathlib import Path
from typing import Dict

from sklearn.model_selection import train_test_split
import yaml

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("logs/feature_extraction.log"),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)


def load_config(config_path: Path) -> Dict:
    """加载并验证配置文件"""
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = yaml.safe_load(f)

        return cfg
    except Exception as e:
        logger.error(f"配置文件加载失败: {str(e)}")
        sys.exit(1)


def validate_files(cfg: Dict):
    """检查图片和标注文件是否匹配"""
    raw_images = set(
        [f for f in os.listdir(cfg["paths"]["raw"]["images"]) if f.lower().endswith((tuple(cfg["dataset"]["image"]["extensions"])))])
    annotations = set([f.replace(".txt", "") for f in os.listdir(cfg["paths"]["raw"]["annotations"]) if f.endswith(".txt")])
    # 排除掉生成的类别文件
    annotations.remove("classes")

    # 找出有效文件对（图片文件名不带后缀与标注文件名匹配）
    valid_files = []
    for img_file in raw_images:
        img_name_without_ext = os.path.splitext(img_file)[0]
        if img_name_without_ext in annotations:
            valid_files.append(img_file)  # 保留带后缀的文件名

    # 打印警告信息
    missing_annotations = [f for f in raw_images
                           if os.path.splitext(f)[0] not in annotations]
    missing_images = set(annotations) - {os.path.splitext(f)[0] for f in raw_images}

    if missing_annotations:
        logger.warn(f"[警告] 缺少标注的图片文件（{len(missing_annotations)}个）: {missing_annotations[:3]}...")
    if missing_images:
        logger.warn(f"[警告] 缺少图片的标注文件（{len(missing_images)}个）: {list(missing_images)[:3]}...")

    assert valid_files, "错误: 没有有效的图片-标注文件对！"
    return valid_files  # 返回带后缀的文件名列表


def create_directory_structure(cfg: Dict):
    """创建YOLOv8所需的目录结构"""
    for dir_type in cfg["paths"]["yolov8_dirs"].values():
        for split_dir in dir_type.values():
            dir_path = os.path.join(cfg["paths"]["processed"]["dir"], split_dir)
            os.makedirs(dir_path, exist_ok=True)
            print(f"创建目录: {dir_path}")


def split_and_copy_data(valid_files, cfg: Dict):
    """划分数据集并复制文件"""
    # 划分训练集和验证集
    train_files, val_files = train_test_split(
        valid_files, test_size=0.2, random_state=cfg["settings"]["random_seed"]
    )

    print(f"\n数据集划分: 训练集={len(train_files)}张, 验证集={len(val_files)}张")

    def copy_files(files, split):
        """复制图片和对应的标注文件"""
        for img_file in files:
            # 源路径（图片）
            src_img = os.path.join(cfg["paths"]["raw"]["images"], img_file)
            dst_img = os.path.join(cfg["paths"]["processed"]["dir"], cfg["paths"]["yolov8_dirs"]["images"][split], img_file)

            # 源路径（标注）
            txt_file = os.path.splitext(img_file)[0] + ".txt"
            src_txt = os.path.join(cfg["paths"]["raw"]["annotations"], txt_file)
            dst_txt = os.path.join(cfg["paths"]["processed"]["dir"], cfg["paths"]["yolov8_dirs"]["labels"][split], txt_file)

            # 执行复制
            try:
                shutil.copy(src_img, dst_img)
                shutil.copy(src_txt, dst_txt)
            except FileNotFoundError as e:
                print(f"错误: 文件复制失败 - {e}")
                continue

    # 执行复制
    print("\n正在复制训练集文件...")
    copy_files(train_files, "train")
    print("正在复制验证集文件...")
    copy_files(val_files, "val")


def create_yaml_config(cfg: Dict):
    """生成YOLOv8的dataset.yaml配置文件"""
    config = {
        "path": cfg["paths"]["processed"]["dir"],  # 使用绝对路径
        "train": "images/train",
        "val": "images/val",
        "names": {i: name for i, name in enumerate(cfg["dataset"]["classes"]["names"])}
    }

    yaml_path = os.path.join(cfg["paths"]["processed"]["dir"], "dataset.yaml")
    with open(yaml_path, "w") as f:
        yaml.dump(config, f, sort_keys=False)

    print(f"\n配置文件已生成: {yaml_path}")
    print("内容预览:")
    print("-" * 30)
    with open(yaml_path, "r") as f:
        print(f.read())
    print("-" * 30)


if __name__ == "__main__":
    logger.info("=" * 50)
    logger.info("YOLOv8 数据预处理（修正版）")
    logger.info("=" * 50 + "\n")

    try:

        # 加载配置
        cfg = load_config(Path("configs/default.yaml"))

        # 步骤1: 检查数据
        logger.info("[1/4] 验证文件匹配性...")
        valid_files = validate_files(cfg)
        logger.info(f"找到有效文件对: {len(valid_files)}个")

        # 步骤2: 创建目录
        logger.info("\n[2/4] 创建目录结构...")
        create_directory_structure(cfg)

        # 步骤3: 处理数据
        logger.info("\n[3/4] 划分并复制数据...")
        split_and_copy_data(valid_files, cfg)

        # 步骤4: 生成配置
        logger.info("\n[4/4] 生成YOLOv8配置...")
        create_yaml_config(cfg)

        logger.info("\n" + "=" * 50)
        logger.info(f'预处理完成！输出目录: {cfg["paths"]["processed"]["dir"]}')
        logger.info("=" * 50)
    except Exception as e:
        logger.error(f"\n[严重错误] 处理失败: {e}")
        logger.error("请检查：")
        logger.error("1. 文件路径是否正确")
        logger.error("2. 图片和标注是否匹配")
        logger.error("3. 文件权限问题")
